// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/functions/sparksql/specialforms/decimal_round.h>
#include <pollux/expression/constant_expr.h>

namespace kumo::pollux::functions::sparksql {
    namespace {
        template<typename TResult, typename TInput>
        class DecimalRoundFunction : public exec::VectorFunction {
        public:
            DecimalRoundFunction(
                int32_t scale,
                uint8_t inputPrecision,
                uint8_t inputScale,
                uint8_t resultPrecision,
                uint8_t resultScale)
                : scale_(
                      scale >= 0
                          ? std::min(
                              scale,
                              static_cast<int32_t>(LongDecimalType::kMaxPrecision))
                          : std::max(
                              scale,
                              -static_cast<int32_t>(LongDecimalType::kMaxPrecision))),
                  inputPrecision_(inputPrecision),
                  inputScale_(inputScale),
                  resultPrecision_(resultPrecision),
                  resultScale_(resultScale) {
                const auto [p, s] = DecimalRoundCallToSpecialForm::getResultPrecisionScale(
                    inputPrecision, inputScale, scale);
                POLLUX_USER_CHECK_EQ(
                    p,
                    resultPrecision,
                    "The result precision of decimal_round is inconsistent with Spark expected.");
                POLLUX_USER_CHECK_EQ(
                    s,
                    resultScale,
                    "The result scale of decimal_round is inconsistent with Spark expected.");

                // Decide the rescale factor of divide and multiply when rounding to a
                // negative scale.
                auto rescaleFactor = [&](int32_t rescale) {
                    POLLUX_USER_CHECK_GT(
                        rescale, 0, "A non-negative rescale value is expected.");
                    return DecimalUtil::kPowersOfTen[std::min(
                        rescale, static_cast<int32_t>(LongDecimalType::kMaxPrecision))];
                };
                if (scale_ < 0) {
                    divideFactor_ = rescaleFactor(inputScale_ - scale_);
                    multiplyFactor_ = rescaleFactor(-scale_);
                }
            }

            void apply(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args,
                const TypePtr &resultType,
                exec::EvalCtx &context,
                VectorPtr &result) const override {
                POLLUX_USER_CHECK(
                    args[0]->is_constant_encoding() || args[0]->is_flat_encoding(),
                    "Single-arg deterministic functions receive their only argument as flat or constant vector.");
                context.ensure_writable(rows, resultType, result);
                result->clear_nulls(rows);
                auto rawResults =
                        result->as_unchecked<FlatVector<TResult> >()->mutableRawValues();
                if (args[0]->is_constant_encoding()) {
                    // Fast path for constant vector.
                    applyConstant(rows, args[0], rawResults);
                } else {
                    // Fast path for flat vector.
                    applyFlat(rows, args[0], rawResults);
                }
            }

            bool supportsFlatNoNullsFastPath() const override {
                return true;
            }

        private:
            inline TResult applyRound(const TInput &input) const {
                if (scale_ >= 0) {
                    TResult rescaledValue;
                    const auto status = DecimalUtil::rescaleWithRoundUp<TInput, TResult>(
                        input,
                        inputPrecision_,
                        inputScale_,
                        resultPrecision_,
                        resultScale_,
                        rescaledValue);
                    POLLUX_DCHECK(status.ok());
                    return rescaledValue;
                } else {
                    TResult rescaledValue;
                    DecimalUtil::divideWithRoundUp<TResult, TInput, int128_t>(
                        rescaledValue, input, divideFactor_.value(), false, 0, 0);
                    rescaledValue *= multiplyFactor_.value();
                    return rescaledValue;
                }
            }

            void applyConstant(
                const SelectivityVector &rows,
                const VectorPtr &arg,
                TResult *rawResults) const {
                const TResult rounded =
                        applyRound(arg->as_unchecked<ConstantVector<TInput> >()->value_at(0));
                rows.applyToSelected([&](auto row) { rawResults[row] = rounded; });
            }

            void applyFlat(
                const SelectivityVector &rows,
                const VectorPtr &arg,
                TResult *rawResults) const {
                auto rawValues = arg->as_unchecked<FlatVector<TInput> >()->mutableRawValues();
                rows.applyToSelected(
                    [&](auto row) { rawResults[row] = applyRound(rawValues[row]); });
            }

            const int32_t scale_;
            const uint8_t inputPrecision_;
            const uint8_t inputScale_;
            const uint8_t resultPrecision_;
            const uint8_t resultScale_;
            std::optional<int128_t> divideFactor_ = std::nullopt;
            std::optional<int128_t> multiplyFactor_ = std::nullopt;
        };

        std::shared_ptr<exec::VectorFunction> createDecimalRound(
            const TypePtr &inputType,
            int32_t scale,
            const TypePtr &resultType) {
            const auto [inputPrecision, inputScale] =
                    getDecimalPrecisionScale(*inputType);
            const auto [resultPrecision, resultScale] =
                    getDecimalPrecisionScale(*resultType);
            if (inputType->isShortDecimal()) {
                if (resultType->isShortDecimal()) {
                    return std::make_shared<DecimalRoundFunction<int64_t, int64_t> >(
                        scale, inputPrecision, inputScale, resultPrecision, resultScale);
                } else {
                    return std::make_shared<DecimalRoundFunction<int128_t, int64_t> >(
                        scale, inputPrecision, inputScale, resultPrecision, resultScale);
                }
            } else {
                if (resultType->isShortDecimal()) {
                    return std::make_shared<DecimalRoundFunction<int64_t, int128_t> >(
                        scale, inputPrecision, inputScale, resultPrecision, resultScale);
                } else {
                    return std::make_shared<DecimalRoundFunction<int128_t, int128_t> >(
                        scale, inputPrecision, inputScale, resultPrecision, resultScale);
                }
            }
        }
    } // namespace

    std::pair<uint8_t, uint8_t>
    DecimalRoundCallToSpecialForm::getResultPrecisionScale(
        uint8_t precision,
        uint8_t scale,
        int32_t roundScale) {
        // After rounding we may need one more digit in the integral part,
        // e.g. 'decimal_round(9.9, 0)' -> '10', 'decimal_round(99, -1)' -> '100'.
        const int32_t integralLeastNumDigits = precision - scale + 1;
        if (roundScale < 0) {
            // Negative scale means we need to adjust `-scale` number of digits before
            // the decimal point, which means we need at least `-scale + 1` digits after
            // rounding, and the result scale is 0.
            const auto newPrecision = std::max(
                integralLeastNumDigits,
                -std::max(
                    roundScale, -static_cast<int32_t>(LongDecimalType::kMaxPrecision)) +
                1);
            // We have to accept the risk of overflow as we can't exceed the max
            // precision.
            return {
                std::min(
                    newPrecision, static_cast<int32_t>(LongDecimalType::kMaxPrecision)),
                0
            };
        }
        const uint8_t newScale = std::min(static_cast<int32_t>(scale), roundScale);
        // We have to accept the risk of overflow as we cannot exceed the max
        // precision.
        return {
            std::min(
                integralLeastNumDigits + newScale,
                static_cast<int32_t>(LongDecimalType::kMaxPrecision)),
            newScale
        };
    }

    TypePtr DecimalRoundCallToSpecialForm::resolveType(
        const std::vector<TypePtr> & /*argTypes*/) {
        POLLUX_FAIL("Decimal round function does not support type resolution.");
    }

    exec::ExprPtr DecimalRoundCallToSpecialForm::constructSpecialForm(
        const TypePtr &type,
        std::vector<exec::ExprPtr> &&args,
        bool trackCpuUsage,
        const core::QueryConfig & /*config*/) {
        POLLUX_USER_CHECK(
            type->isDecimal(),
            "The result type of decimal_round should be decimal type.");
        POLLUX_USER_CHECK_GE(
            args.size(), 1, "Decimal_round expects one or two arguments.");
        POLLUX_USER_CHECK_LE(
            args.size(), 2, "Decimal_round expects one or two arguments.");
        POLLUX_USER_CHECK(
            args[0]->type()->isDecimal(),
            "The first argument of decimal_round should be of decimal type.");

        int32_t scale = 0;
        if (args.size() > 1) {
            POLLUX_USER_CHECK_EQ(
                args[1]->type()->kind(),
                TypeKind::INTEGER,
                "The second argument of decimal_round should be of integer type.");
            auto constantExpr = std::dynamic_pointer_cast<exec::ConstantExpr>(args[1]);
            POLLUX_USER_CHECK_NOT_NULL(
                constantExpr,
                "The second argument of decimal_round should be constant expression.");
            POLLUX_USER_CHECK(
                constantExpr->value()->is_constant_encoding(),
                "The second argument of decimal_round should be wrapped in constant vector.");
            auto constant_vector =
                    constantExpr->value()->as_unchecked<ConstantVector<int32_t> >();
            POLLUX_USER_CHECK(
                !constant_vector->is_null_at(0),
                "The second argument of decimal_round is non-nullable.");
            scale = constant_vector->value_at(0);
        }

        auto decimalRound = createDecimalRound(args[0]->type(), scale, type);

        return std::make_shared<exec::Expr>(
            type,
            std::move(args),
            std::move(decimalRound),
            exec::VectorFunctionMetadata{},
            kRoundDecimal,
            trackCpuUsage);
    }
} // namespace kumo::pollux::functions::sparksql
